#include "StdAfx.h"
#include "EventMonWbem.h"

CEventMonWbem::CEventMonWbem()
{
	m_pSvc = NULL;
	m_pLoc = NULL;

	m_EvtObj = NULL;
}

CEventMonWbem::~CEventMonWbem(void)
{
	if(m_pSvc != NULL)
	{
		m_pSvc->Release();
	}

	if(m_pLoc != NULL)
	{
		m_pLoc->Release();
	}

	CoUninitialize();
}

//
// Initializes the Wbem Object and adjusts the security.
//
bool 
CEventMonWbem::Init(void)
{
    HRESULT hr;

	// Adjust security to give rights to monitor eventlog.
	if(!AdjustProcessSecurity())
	{
		return false;
	}

	// Setup connection to Wbem Object
	CoInitialize(NULL);

    hr =  CoInitializeSecurity(
        NULL, 
        -1,                          // COM negotiates service
        NULL,                        // Authentication services
        NULL,                        // Reserved
        RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication 
        RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation  
        NULL,                        // Authentication info
        EOAC_NONE,                   // Additional capabilities 
        NULL                         // Reserved
        );

    if (FAILED(hr))
    {
        return false;     // Program has failed.
    }

	hr = CoCreateInstance(CLSID_WbemLocator, 0, 
        CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID *) &m_pLoc);
 
    if (FAILED(hr))
    {
        return false;     // Program has failed.
    }

	// Connect to the root\default namespace with the current user.
    hr = m_pLoc->ConnectServer(
            BSTR(L"root\\cimv2"), 
            NULL, NULL, 0, NULL, 0, 0, &m_pSvc);

    if (FAILED(hr))
    {
		m_pLoc->Release();
        return false;     // Program has failed.
    }

	// Set the proxy so that impersonation of the client occurs.
    hr = CoSetProxyBlanket(m_pSvc,
       RPC_C_AUTHN_WINNT,
       RPC_C_AUTHZ_NONE,
       NULL,
       RPC_C_AUTHN_LEVEL_DEFAULT,
       RPC_C_IMP_LEVEL_IMPERSONATE,
       NULL,
       EOAC_NONE
    );

    if (FAILED(hr))
    {
        return false;      // Program has failed.
    }

	return true;
}

//
// Adjusts the security token for SeSecurityPrivilege so that we can read
// the event log.
//
bool 
CEventMonWbem::AdjustProcessSecurity(void)
{
	HANDLE Token;

	if(!OpenProcessToken(GetCurrentProcess() ,TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &Token)){
		return false;
	} else {
		TOKEN_PRIVILEGES *TokPriv;
		DWORD retlen = 0;

		// Call first with a NULL buffer to determine the size to allocate.
		BOOL Result =  GetTokenInformation( Token,				// TokenHandle
											TokenPrivileges,	// TokenInformationClass
											NULL,				// TokenInformation
											retlen,				// TokenInformationLength
											&retlen				// ReturnLength
											);

		if((TokPriv = ((struct _TOKEN_PRIVILEGES *)malloc(retlen))) == NULL){
			return false;
		}

		Result =  GetTokenInformation( Token,					// TokenHandle
											TokenPrivileges,	// TokenInformationClass
											TokPriv,			// TokenInformation
											retlen,				// TokenInformationLength
											&retlen				// ReturnLength
											);

		if(FAILED(Result)){
			free(TokPriv);
			return false;
		}

		// Find the specific LUID to locate via the friendly name.
		LUID SecurityLUID; 
		LookupPrivilegeValue(NULL, L"SeSecurityPrivilege", &SecurityLUID);

		// Find the match and enable this previlege
		for(DWORD n=0;n < TokPriv->PrivilegeCount; n++)
		{
			if(TokPriv->Privileges[n].Luid.HighPart == SecurityLUID.HighPart && 
				TokPriv->Privileges[n].Luid.LowPart == SecurityLUID.LowPart)
			{
				TokPriv->Privileges[n].Attributes = SE_PRIVILEGE_ENABLED;
			} else {
				TokPriv->Privileges[n].Attributes = SE_PRIVILEGE_ENABLED;
			}
		}

		// Now adjust the new token with the changed privilege.
		Result = AdjustTokenPrivileges(Token, FALSE, TokPriv, retlen, NULL, &retlen);
		
		if(FAILED(Result))
		{
			free(TokPriv);
			return false;
		}
	}

	return true;
}

//
// Creates the Monitor Event that will be used to wait for the event number
// in question.
//
HRESULT 
CEventMonWbem::MonitorEvt(WCHAR * m_EventNum, WCHAR * EventLog)
{
	WCHAR QueryStr[1024];

	// Create the query string.  If EventLog is not NULL use it as well as part of the query.
	if(EventLog == NULL){
		wsprintf(QueryStr, L"Select * from __InstanceCreationEvent Where TargetInstance ISA 'Win32_NTLogEvent' and TargetInstance.EventCode = '%s'", m_EventNum);
	} else {
		wsprintf(QueryStr, L"Select * from __InstanceCreationEvent Where TargetInstance ISA 'Win32_NTLogEvent' and TargetInstance.EventCode = '%s' and (TargetInstance.m_LogFile = '%s' or TargetInstance.SourceName = '%s')", m_EventNum, EventLog, EventLog);
	}

    // Creates an EventQuery::Indicate method to be used to wait for the next matching result.
	HRESULT hRes;
    hRes = m_pSvc->ExecNotificationQuery(
		L"WQL", 
        QueryStr,
        WBEM_FLAG_RETURN_IMMEDIATELY | WBEM_FLAG_FORWARD_ONLY, 
        NULL, 
        &m_EvtObj);

	return hRes;
}

//
// Uses the Monitor Event created by MonitorEvt to wait for
// an event to be fired.
//
HRESULT 
CEventMonWbem::WaitForEvt(long TimeOut, ULONG * NumEvents)
{
	IWbemClassObject* apObj[1];

	// Wait for 1 second.  Return the result from the query.  This could
	// be WBEM_S_NO_ERROR or WBEM_S_TIMEDOUT.  But NumEvents > 0 is the
	// indication that an event was returned.
	HRESULT result = m_EvtObj->Next(1000, 1, apObj, NumEvents);

	for(ULONG n=0; n < *NumEvents; n++)
	{
		apObj[n]->Release();
	}

	return result;
}
